import torch
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

# torchvision数据集输出PILImage格式
# 改变数据区间从[0,1]==>[-1,1]张量格式
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# 训练集
# windows环境下num_workers=0
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=0)

# 测试集
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=0)

# 标签
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')


def imshow(img):
    img = img / 2 + 0.5
    # img是tensor类型
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


# 从数据迭代器读取一张图片
dataiter = iter(trainloader)
images, labels = dataiter.next()


# labels返回classes对应的index
# print(labels)

# 展示图片
# imshow(torchvision.utils.make_grid(images))
# 打印label
# print(' '.join('%5s' % classes[labels[j]] for j in range(4)))

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 卷积层
        # 卷积核5*5
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        # 池化层
        self.pool = nn.MaxPool2d(2, 2)
        # 全连接层
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 64)
        self.fc3 = nn.Linear(64, 10)

    def forward(self, x):
        # 任意卷积层后需要加上 激活层relu 和 池化层max_pool2d
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))

        # print(x.size())
        # torch.Size([1, 16, 6, 6])

        # 卷积处理后，需要调整张量的形状
        x = x.view(-1, 16 * 5 * 5)

        # 全连接
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


net = Net()
print(net)

# 损失函数定义
# 交叉熵
criterion = nn.CrossEntropyLoss()
# 随机梯度下降优化器，parameters可训练参数，动量0.9，学习率0.01
optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9)

# 训练模型
for epoch in range(2):
    running_loss = 0.0
    for i, data in enumerate(trainloader):
        # 从data中取出图像张量inputs 标签张量labels
        inputs, labels = data
        # 梯度清零
        optimizer.zero_grad()
        # 图像输入网络
        outputs = net(inputs)
        # 计算损失
        loss = criterion(outputs, labels)
        # 反向传播
        loss.backward()
        # 梯度更新
        optimizer.step()
        # 打印训练信息
        running_loss += loss.item()
        if (i + 1) % 2000 == 0:
            print('[%d,%5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
            running_loss = 0

print('finish')
PATH = './CIFAR_NET.pth'
torch.save(net.state_dict(), PATH)
